import os
import subprocess
import mindspore
from mindspore import nn, context
from thyassist.machine_learning.dataloader.load_gan_data import create_dataset, load_ckpt, ImagePool, get_lr, Reporter
from thyassist.machine_learning.loss import DiscriminatorLoss, GeneratorLoss
from thyassist.machine_learning.networks import get_generator, get_discriminator, Generator, TrainOneStepG, TrainOneStepD
from launcher import get_project_root

mindspore.set_seed(1)
download_dir = get_project_root()
USE_ORANGE_PI = False
mindspore.set_seed(1)

if os.name == 'nt':
    # Windows系统下采用CPU环境
    context.set_context(mode=mindspore.PYNATIVE_MODE, pynative_synchronize=True)
    mindspore.set_device(device_target='CPU')
else:
    try:  # 先尝试昇腾环境，若不可用则尝试GPU，仍不可用则使用CPU
        if subprocess.run(['whoami'], capture_output=True, text=True, check=True).stdout.strip() == 'HwHiAiUser':
            context.set_context(mode=context.GRAPH_MODE, jit_config={"jit_level": "O2"})
            mindspore.set_device(device_target='Ascend')
            USE_ORANGE_PI = True
        else:
            context.set_context(mode=context.GRAPH_MODE)
            mindspore.set_device(device_target='Ascend')
    except:
        try:
            context.set_context(mode=context.GRAPH_MODE)
            mindspore.set_device(device_target="GPU")
        except:
            context.set_context(mode=context.GRAPH_MODE)
            mindspore.set_device(device_target="CPU")


def train(train_data_path, is_load_ckpt, rank=0, num_epoch=200, train_batch_size=16, image_size=512):
    ds = create_dataset(data_path=train_data_path, is_train=True, batch_size=train_batch_size, device_num=1, shuffle=False, image_size=image_size)
    step_size = ds.get_dataset_size()

    G_A = get_generator(image_size=image_size, gl_num=8)
    G_B = get_generator(image_size=image_size, gl_num=8)
    D_A = get_discriminator()
    D_B = get_discriminator()
    if is_load_ckpt:
        load_ckpt(G_A=G_A, G_B=G_B, D_A=D_A, D_B=D_B, is_load_DA=True, is_load_DB=True, is_load_GA=True, is_load_GB=True)
    imgae_pool_A = ImagePool(50)
    imgae_pool_B = ImagePool(50)
    generator = Generator(G_A, G_B)

    loss_D = DiscriminatorLoss(D_A, D_B)
    loss_G = GeneratorLoss(generator, D_A, D_B)
    optimizer_G = nn.Adam(generator.trainable_params(),
                          get_lr(lr_init=0.0002,
                                 num_epoch=200,
                                 warm_up_epoch=100,
                                 n_epochs_decay=100,
                                 step_size=step_size),
                                 beta1=0.5)
    optimizer_D = nn.Adam(loss_D.trainable_params(),
                          get_lr(lr_init=0.0002,
                                 num_epoch=200,
                                 warm_up_epoch=100,
                                 n_epochs_decay=100,
                                 step_size=step_size),
                          beta1=0.5)

    net_G = TrainOneStepG(loss_G, generator, optimizer_G)
    net_D = TrainOneStepD(loss_D, optimizer_D)

    data_loader = ds.create_dict_iterator()
    if rank == 0:
        reporter = Reporter(step_size=step_size, is_save_checkpoint=False)
        reporter.info('==========start training===============')
    for _ in range(num_epoch):
        if rank == 0:
            reporter.epoch_start()
        for data in data_loader:
            img_A = data["image_A"]
            img_B = data["image_B"]
            res_G = net_G(img_A, img_B)
            fake_A = res_G[0]
            fake_B = res_G[1]
            res_D = net_D(img_A, img_B, imgae_pool_A.query(fake_A), imgae_pool_B.query(fake_B))
            if rank == 0:
                reporter.step_end(res_G, res_D)
                reporter.visualizer(img_A, img_B, fake_A, fake_B)
        if rank == 0:
            reporter.epoch_end(net_G)

    mindspore.save_checkpoint(net_G.G.generator.G_A, os.path.join(get_project_root(),  "gan_checkpoints", "g_a"))
    mindspore.save_checkpoint(net_G.G.generator.G_B, os.path.join(get_project_root(), "gan_checkpoints", "g_b"))
    mindspore.save_checkpoint(net_G.G.D_A, os.path.join(get_project_root(), "gan_checkpoints", "d_a.ckpt"))
    mindspore.save_checkpoint(net_G.G.D_B, os.path.join(get_project_root(), "gan_checkpoints", "d_b.ckpt"))
    if rank == 0:
        reporter.info('==========end training===============')


if __name__ == "__main__":
    train(train_data_path=os.path.join(get_project_root(), "gan_datasets"), is_load_ckpt=False)
